package com.interview.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.FileSystemResource;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.util.List;

@Tag(name = "RagDemoController", description = "RagDemoController")
@Slf4j
@RestController
@RequestMapping("/rag")
@AllArgsConstructor
public class RagDemoController {
    @Autowired
    private VectorStore vectorStore;

    @Operation(summary = "模型训练")
    @PostMapping("/add")
    public String add(@RequestParam("file") MultipartFile file) throws IOException {
        File tmpFile = File.createTempFile("upload-", file.getOriginalFilename());
        file.transferTo(tmpFile);
        addVectorStore(tmpFile);
        tmpFile.deleteOnExit();
        return "文件已成功添加到向量库";
    }

    public void addVectorStore(File documentation) {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(new FileSystemResource(documentation));
        TokenTextSplitter splitter = new TokenTextSplitter(300, 200, 10, 400, true);
        List<Document> documents = splitter.apply(tikaDocumentReader.get());
        documents.forEach(document -> {
            vectorStore.add(List.of(
                    new Document(document.getFormattedContent())));
        });
    }
}